#!/usr/bin/env python3
'''
已验证，验证通过
验证结果
1. 对比之前飞船会主动追击、预判敌人
2. 会主动的左右攻击
3. 关卡到后期会出现新的游戏模式，可能是因为这个导致后续分数提交不高
todo
1. 对比之前的代码，总结是因为做了哪些更改才提升了分数
2. 调整代码，使得模型可以继续训练，直到通关（是因为后续学习率不高导致的新环境无法调整，还是卷积的参数量太少导致无法记录有效数据趋势）

验证结果：
1. 可以训练，但是发现训练到中途奖励达到一定的值后，就会下降，然后测试奖励回报一直都是0，不清楚原因是什么
2. 通过测试最好的保存模型可知，确实是经过训练的，因为之前的关卡只有上面的敌人不存在左右两边的敌人，所以经过训练后，不会随意的向左右两边开炮，只会向上开炮消灭敌人，而且会守株待兔，不随便移动击杀敌人，基本上可以维持击杀31个敌人记录。但是由于更后面的关卡敌人会从左右两边攻击，由于模型此时未训练应对这个场景，所以会被敌人击杀
3. 并且经过训练的模型可以在前面几关基本可以不丢命

目前遇到的问题就是，为什么训练到后面奖励会下降，测试奖励一直都是0，不清楚原因是什么，需要解决这个问题可能才能进一步训练模型

另外发现训练到后面有左右两边敌人的时候，即使还剩3条命，游戏依旧会结束，这和一开始的即使被击杀只要有剩命就会继续游戏不一样，不清楚是什么原因导致的，需要进一步研究（这个问题不用研究了，因为是EpisodicLifeEnv包装器导致即使还有命也会当作游戏结束处理，不知道是不是因为这个原因导致游戏训练到后期就无法继续训练下去了）


为了进一步提高分数，开始研究如何调整

分数进一步提高了，可以达到1000+分，平均分也达到了600+，可能是之前的提交产生了作用（增加了学习率优化器）
不过由于是训练中途增加了学习率迭代器，大概是在2600000（也就是在4万多次训练后），增加了迭代器，当时增加时发现分数一会高一会低，所以才增加了学习率迭代器，
测试降低学习率试试，并且这个时候对已经训练的模型进行了一次中途保存，用于后续测试是否需要去除迭代其试试
分数达到了2000+分，此时经过计算加入了学习率衰减迭代了大概2~4次
大概是在轮数2600000到7064192之间增加了学习率迭代器是移植在上升的，之后就是一直处于不断的震荡，考虑是否需要搜东设置恢复学习率到某个值

更新了探索熵，发现回报奖励震荡减轻，可以出现较多的800+高分，测试回报奖励也重新出现了2000+分，感觉熵也会导致分数震荡，可能不是学习率的影响（因为前期需要探索，后期就需要使用自己学习到的知识进行游戏的训练，这也是其他人的代码会做的事情，gpt也提供了相关的逐渐降低熵的公式），考虑是否需要将学习率恢复到4e-5左右的值试试

更新训练记录，进一步下降探索熵到0.001以及0.0001，减少学习率的更新频率（到了1.6e-5左右），训练平均分数可以达到900+分左右（但是还是会出现震荡下降到700+分），测试环境的分数
依旧可以达到2000+分。没有明显的提升，不清楚是探索熵的问题、还是学习率的问题，还是模型的拟合能力不够强，需要进一步研究

20241215: 正在重新调整网络结构
20241216：完成调整，重新训练，再2号机上训练，训练分数50分，测试分数340分，继续训练
20241217: 学习率0.0003645，训练分数95分，测试分数353分，继续训练
20241218：学习率 0.00023914845,训练分数104分，测试分数382分，继续训练
20241219: 学习率 0.0001569，训练分数144分，测试分数390分，确定学习率是否下降太快，基本上一天下降3次，学习率下降无问题，继续训练
20241220：学习率 0.0001029455,训练分数185分，测试分数406分，可以看出模型的拟合能力有所提高，继续训练
20241221:学习率 6.754258588364964e-05,训练分数185分，测试分数412分，继续训练
20241223:学习率  4.923854510918059e-05,训练分数178分，测试分数441分，继续训练
20241224：学习率3.589489938459265e-05，训练记录211分，测试分数452分，继续训练
20241225：学习率 2.119557913760812e-05,训练分数218分，测试分数452分，继续训练
20241226：学习率1.545157719131632e-05，训练分数228分，测试分数700分，继续训练
20241227：学习率1.0137779795222638e-05，训练分数241分，测试分数700分，继续训练
20241228：学习率6.651397323645573e-06，训练分数246分，测试分数700分，继续训练
20241229:学习率 3.927583605639475e-06,训练分数251分，测试分数700分，继续训练
20241230:学习率2.57688760366006e-06，训练分数246分，测试分数700分（新分691分），继续训练
20241231:学习率1.6906959567613658e-06，训练分数240分，测试分数700分（新分692），play模型，停止训练
'''
import gymnasium as gym
import ptan
import numpy as np
import argparse
from tensorboardX import SummaryWriter
import os

import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim

from lib import common
import ale_py

gym.register_envs(ale_py)
GAMMA = 0.99
LEARNING_RATE = 5e-4
ENTROPY_BETA = 0.01
# ENTROPY_BETA = 0.02 # 降低熵的权重，看看效果
NUM_ENVS = 16

REWARD_STEPS = 4
CLIP_GRAD = 0.5

SAVE_ITERS = 100


class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(input_shape)
        self.policy = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 255.0
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)


def unpack_batch(batch, net, device="cpu"):
    """
    Convert batch into training tensors
    :param batch:
    :param net:
    :return: states variable, actions tensor, reference values variable
    """
    states = []
    actions = []
    rewards = []
    not_done_idx = []
    last_states = []
    for idx, exp in enumerate(batch):
        states.append(np.array(exp.state, copy=False))
        actions.append(int(exp.action))
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(np.array(exp.last_state, copy=False))
    states_v = torch.FloatTensor(states).to(device)
    actions_t = torch.LongTensor(actions).to(device)

    # handle rewards
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        last_states_v = torch.FloatTensor(last_states).to(device)
        last_vals_v = net(last_states_v)[1]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_t, ref_vals_v


def set_seed(seed, envs=None, cuda=False):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda:
        torch.cuda.manual_seed(seed)

    if envs:
        for idx, env in enumerate(envs):
            env.reset(seed=seed + idx)


'''
测试来看，基础得分可以是21分起，后续评估是否需要裁剪奖励
'''

class RewardWrapper(gym.Wrapper):
    def __init__(self, env):
        super(RewardWrapper, self).__init__(env)
        self.previous_lives = None  # 用于跟踪智能体的生命值

    def reset(self, **kwargs):
        # 重置环境并初始化生命值
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info['lives']
        return obs, info

    def step(self, action):
        obs, reward, done, trunc, info = self.env.step(action)

        # 自定义奖励机制
        custom_reward = self.custom_reward(info, reward, done, action)

        return obs, custom_reward, done, trunc, info

    def custom_reward(self, info, original_reward, done, action):
        custom_reward = 0

        # # 无效射击惩罚 去除无效的射击惩罚，因为执行动作后需要一定的时间才能看到效果，所以不会立即反馈
        # if (action == 1 or action == 5 or action == 6) and original_reward < 1:
        #     original_reward -= 1

        # 基于原始游戏的奖励，如击杀外星人时的正向奖励
        custom_reward += original_reward
        if custom_reward != 0:
            custom_reward //= 21
            if custom_reward == 0:
                # 低于21分的奖励都是10分，已记录
                # print(f"检测到低于21分的奖励，进行记录：{original_reward}")
                custom_reward = 1

        # 获取当前生命值
        current_lives = info['lives']

        # 检查是否死亡并给予惩罚
        if current_lives < self.previous_lives:
            custom_reward -= 10  # 死亡惩罚
        elif current_lives > self.previous_lives:
            custom_reward += 10
        self.previous_lives = current_lives

        # 存活奖励（每一步给予小奖励）
        # if not done:
        #     custom_reward += 0.1  # 每步存活的奖励

        # 如果有其他道具收集机制，你可以在这里添加道具收集奖励
        # if self.item_collected():
        #     custom_reward += 5  # 收集道具奖励

        return custom_reward


def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    # if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        # env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=3)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = ptan.common.wrappers.FireResetEnv(env)
    env = ptan.common.wrappers.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = RewardWrapper(env)
    # if reward_clipping:
        # env = ptan.common.wrappers.ClippedRewardsWrapper(env)
    return env

def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v, _ = net(obs_v)
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes

def select_device(args):
    if args.cuda and torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available() and args.cuda:
        return torch.device("mps")
    return torch.device("cpu")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    args = parser.parse_args()
    device = select_device(args=args)

    make_env = lambda: wrap_dqn(gym.make("ALE/Assault-v5", obs_type='rgb', frameskip=4, repeat_action_probability=0.0))
    envs = [make_env() for _ in range(NUM_ENVS)]
    test_env = make_env()
    writer = SummaryWriter(comment="-assault-a2c-rollouts_" + args.name)
    set_seed(20, envs, cuda=args.cuda)

    net = AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    print(net)

    # 不同点2
    agent = ptan.agent.ActorCriticAgent(net, apply_softmax=True, device=device)
    # 不同点3 todo
    exp_source = ptan.experience.ExperienceSourceRollouts(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)

    # optimizer = optim.RMSprop(net.parameters(), lr=LEARNING_RATE, eps=1e-5)
    # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)
    # todo 切换一下学习率优化器试试
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)


    # 创建保存训练过程中数据的路径
    save_path = os.path.join("saves", "a2c-" + args.name)
    os.makedirs(save_path, exist_ok=True)
    step_idx = 0
    train_count = 0

    # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        step_idx = checkpoint['step_idx']
        train_count = checkpoint['train_count']
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        print(f"Current learning rate: {scheduler.get_last_lr()[0]}")
        print("加载模型成功")

    best_reward = -float('inf')
    with common.RewardTracker(writer, stop_reward=999999) as tracker:
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            # mb_states: 游戏的状态
            # mb_rewards: 这里应该基于动作获取的实际奖励并累计了未来奖励的计算后得到的Q值
            # mb_actions: 执行的动作
            # mb_values: 对于执行动作前的游戏状态的评价值（Q值），在这份代码中，预测的Q值是需要考虑到未来多步的累计降级
            for mb_states, mb_rewards, mb_actions, mb_values in exp_source:
                # handle new rewards
                # 计算采集游戏数据的奖励，并记录到tensorboard中
                # 为什么这里会没有使用类似缓冲区的方式记录游戏数据
                # 因为这里采用了记录并返回连续游戏轨迹的方式，所以这里不需要缓冲区
                # 因为直接拿到的数据就是一个连续的游戏轨迹数据，可以直接用于训练
                #  类似pong_a2c
                new_rewards = exp_source.pop_total_rewards()
                if new_rewards:
                    if tracker.reward(np.mean(new_rewards), step_idx):
                        break

                optimizer.zero_grad()
                states_v = torch.FloatTensor(mb_states).to(device)
                # 计算动作优势
                mb_adv = mb_rewards - mb_values
                adv_v = torch.FloatTensor(mb_adv).to(device)
                actions_t = torch.LongTensor(mb_actions).to(device)
                vals_ref_v = torch.FloatTensor(mb_rewards).to(device)

                logits_v, value_v = net(states_v)
                #  vals_ref_v是已经转换为Q值的奖励了吗？
                #  没错，根据代码来看，是因为转换为了类似Q值的考虑到未来N步的折扣奖励
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                # 计算动作优势值
                log_prob_v = F.log_softmax(logits_v, dim=1)
                log_prob_actions_v = adv_v * log_prob_v[range(len(mb_states)), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                # 计算动作熵
                prob_v = F.softmax(logits_v, dim=1)
                entropy_loss_v = (prob_v * log_prob_v).sum(dim=1).mean()

                # apply entropy and value gradients
                # 计算损失优化
                loss_v = loss_policy_v + ENTROPY_BETA * entropy_loss_v + loss_value_v
                loss_v.backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                optimizer.step()

                tb_tracker.track("advantage",       adv_v, step_idx)
                tb_tracker.track("values",          value_v, step_idx)
                tb_tracker.track("batch_rewards",   vals_ref_v, step_idx)
                tb_tracker.track("loss_entropy",    entropy_loss_v, step_idx)
                tb_tracker.track("loss_policy",     loss_policy_v, step_idx)
                tb_tracker.track("loss_value",      loss_value_v, step_idx)
                tb_tracker.track("loss_total",      loss_v, step_idx)
                step_idx += NUM_ENVS * REWARD_STEPS
                train_count += 1
                scheduler.step()

                if train_count % 50 == 0:
                    checkpoint = {
                        "net": net.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "step_idx": step_idx,
                        'scheduler_state_dict': scheduler.state_dict(),
                        'train_count': train_count
                    }
                    common.save_checkpoints(step_idx, checkpoint, save_path, "a2c", keep_last=5)
                
                if train_count % 100 == 0:
                    # Test the model
                    test_reward = test_model(test_env, net, device=device, episodes=1)
                    print(f"Test reward: {test_reward:.2f}")

                    # Save the model if it is the best so far
                    if test_reward > best_reward:
                        best_reward = test_reward
                        print(f"New best model saved with reward {best_reward:.2f}")
                    common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-best", keep_best=10)
                    
